// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gtest/gtest.h>

#include <pollux/exec/aggregate.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>
#include <pollux/functions/prestosql/aggregates/register_aggregate_functions.h>

namespace kumo::pollux::aggregate::test {

class AggregationFunctionRegTest : public testing::Test {
 protected:
  void clearAndCheckRegistry() {
    exec::aggregateFunctions().withWLock([&](auto& functionsMap) {
      functionsMap.clear();
      EXPECT_EQ(0, functionsMap.size());
    });
  }
};

TEST_F(AggregationFunctionRegTest, prefix) {
  // Remove all functions and check for no entries.
  clearAndCheckRegistry();

  // Register without prefix and memorize function maps.
  aggregate::prestosql::registerAllAggregateFunctions();
  const auto aggrFuncMapBase = exec::aggregateFunctions().copy();

  // Remove all functions and check for no entries.
  clearAndCheckRegistry();

  // Register with prefix and check all functions have the prefix.
  const std::string prefix{"test.abc_schema."};
  aggregate::prestosql::registerAllAggregateFunctions(prefix);
  exec::aggregateFunctions().withRLock([&](const auto& aggrFuncMap) {
    for (const auto& entry : aggrFuncMap) {
      EXPECT_EQ(prefix, entry.first.substr(0, prefix.size()));
      EXPECT_EQ(1, aggrFuncMapBase.count(entry.first.substr(prefix.size())));
    }
  });
}

TEST_F(AggregationFunctionRegTest, orderSensitive) {
  // Remove all functions and check for no entries.
  clearAndCheckRegistry();

  std::set<std::string> nonOrderSensitiveFunctions = {
      "sum",
      "avg",
      "min",
      "max",
      "count",
      "bool_and",
      "bool_or",
      "bitwise_and_agg",
      "bitwise_or_agg",
      "every",
      "checksum",
      "count_if",
      "geometric_mean",
      "histogram",
      "reduce_agg"};
  aggregate::prestosql::registerAllAggregateFunctions();
  for (const auto& entry : nonOrderSensitiveFunctions) {
    ASSERT_FALSE(exec::getAggregateFunctionMetadata(entry).orderSensitive);
  }

  // Test some but not all order sensitive functions
  std::set<std::string> orderSensitiveFunctions = {
      "array_agg", "arbitrary", "any_value", "map_agg", "map_union", "set_agg"};
  for (const auto& entry : orderSensitiveFunctions) {
    ASSERT_TRUE(exec::getAggregateFunctionMetadata(entry).orderSensitive);
  }
}

TEST_F(AggregationFunctionRegTest, prestoSupportedSignatures) {
  // Remove all functions and check for no entries.
  clearAndCheckRegistry();

  // Register without prefix and all signatures
  aggregate::prestosql::registerAllAggregateFunctions("", true, false);
  auto aggrFuncMapBase = exec::aggregateFunctions().copy();
  EXPECT_EQ(aggrFuncMapBase[kBitwiseAnd].signatures.size(), 4);
  EXPECT_EQ(aggrFuncMapBase[kBitwiseOr].signatures.size(), 4);
  EXPECT_EQ(aggrFuncMapBase[kBitwiseXor].signatures.size(), 4);

  // Register without prefix and only those signatures that are supported by
  // Presto.
  clearAndCheckRegistry();
  aggregate::prestosql::registerAllAggregateFunctions("", true, true);
  aggrFuncMapBase = exec::aggregateFunctions().copy();
  EXPECT_EQ(aggrFuncMapBase[kBitwiseAnd].signatures.size(), 1);
  EXPECT_EQ(aggrFuncMapBase[kBitwiseOr].signatures.size(), 1);
  EXPECT_EQ(aggrFuncMapBase[kBitwiseXor].signatures.size(), 1);

  // Revert to previous state.
  clearAndCheckRegistry();
}

TEST_F(AggregationFunctionRegTest, companionFunction) {
  // Remove all functions and check for no entries.
  clearAndCheckRegistry();

  aggregate::prestosql::registerAllAggregateFunctions();
  const auto aggregates = {"approx_distinct", "count", "sum"};
  const auto companionFunctions = {
      "approx_distinct_merge", "approx_distinct_partial"};

  for (const auto& function : aggregates) {
    ASSERT_FALSE(
        exec::getAggregateFunctionMetadata(function).companionFunction);
  }
  for (const auto& function : companionFunctions) {
    ASSERT_TRUE(exec::getAggregateFunctionMetadata(function).companionFunction);
  }
}

} // namespace kumo::pollux::aggregate::test
